# -*- coding: utf-8 -*-
"""
Given a filename, acquisition time and threshold, it gives the detector number
gain shift, and answer as to whether or not the gain shift has exceeded the 
threshold.

Usage:
    python PBAR_Zspec_Checkain.py <full filename of zspec csv file> <acquisition time> <threshold>

Example:
    python PBAR_Zspec_Checkain.py ~/vxWorks/analysis/ea68.csv 300.0 0.05
 - aquisition time is 300 seconds
 - threshold is 5% (bad if gain shift > 5%)

Created on Tue Aug 27 10:43:08 2013

@author: jkwong
"""

#import os
import sys
import numpy as np

def ReadZspec(fullfilenameList):
    """ Read Zspec data """
    
    # Check if the input is a string (single file)
    if isinstance(fullfilenameList, str):
        temp = []
        temp.append(fullfilenameList)
        fullfilenameList = temp
    
    dat = list()
    for ii in range(len(fullfilenameList)):
        temp = np.genfromtxt(fullfilenameList[ii], \
                                 delimiter=',', \
                                 skip_header = 0, \
                                 skip_footer = 0, \
                                 dtype = 'uint32')
        if (temp.shape[0] == 257):  # shave off the table labels if present
            temp = temp[1:,:]
        if (temp.shape[1] == 138):
            temp = temp[:,1:-1]
        if (temp.shape[1] == 137): # something with newline causeing it to think there is one more detetor
            temp = temp[:,1:]
        dat.append(temp)
    return dat

def CheckGain(*arg):
    """Given a filename, acquisition time and threshold, it gives the detector number
        gain shift, and answer as to whether or not the gain shift has exceeded the 
        threshold.
        8/28/2013, John Kwong
        8/29/2013, Added option to print all or 
    """
    #parse the arguments
    fullFilename, acqTime, threshold = arg[0:3]
    # convert the threshold string to float
    threshold = float(threshold)
    # convert acquisition time
    acqTime = float(acqTime)
    if len(arg) < 4:
        option = 'all'
    else:
        option = arg[3]
    
    # check if the option is valid
    if option not in ['badgain', 'gooddet', 'all', 'badgaingooddet']:
        print("Invalid Option.")
        return

    # Baseline values
    # Fit exponential to this region
    rateRange = np.array([0.5, 2.0])

    # bin value of the fitted exponential at count rate
    rateLook = 1.0    
    # should be at this value
    binBaseline = 100.0
    # the bin value have to be between these values
    binBounds = [20, 200]
    
    # index starts at 1
    badDetectors = np.array([1,2,3,4,5,6,7,8,20,26,31,33,34,38,39,40,44,53,56,62,68,76,80,125,126,127,128,129,130,131,132,133,134,135,136])
        
    # Read the file
    dat = ReadZspec(fullFilename)
    
    # Shift among PMTs - OLD STUFF 8/28/2013
#    binMeanCC = np.zeros(dat[0].shape[1])
#    gainShift = np.zeros(dat[0].shape[1])
    
#    print dat[0].shape
    binArray = np.arange(dat[0].shape[0])
    
    # fit an exponential
#    pfitAll = np.zeros((dat[0].shape[1], 2))
    binMeanAll = np.zeros(dat[0].shape[1])

    for i in xrange(dat[0].shape[1]):
        cutt = (binArray > binBounds[0]) & (binArray < binBounds[1]) & (dat[0][:,i]/acqTime > rateRange[0]) & (dat[0][:,i]/acqTime < rateRange[1])
        if any(cutt):
            temp = np.polyfit(binArray[cutt], np.log(dat[0][cutt,i]/acqTime),1)
            pfit0 = [temp[0], np.exp(temp[1])]
            binMean = (1.0/pfit0[0]) * np.log(rateLook/pfit0[1])
        else:
            binMean = 0.0
        binMeanAll[i] = binMean
#        print i+1, binMean
    gainShift = binMeanAll / binBaseline

#    # OLD STUFF - 8/28/2013, JK
#    binMat = np.tile(np.arange(dat[0].shape[0]),(dat[0].shape[1],1)).T
#    cutt = (binMat > binBounds[0]) & (binMat < binBounds[1]) & (dat[0]/acqTime > rateRange[0]) & (dat[0]/acqTime < rateRange[1])
#
#    # calculate the bin value for all detectors
#    binMeanCC = (binMat * dat[0] * cutt).sum(axis = 0).astype(float) / (dat[0] * cutt).sum(axis = 0).astype(float)
#
#    gainShift = binMeanCC / np.tile(binBaseline, dat[0].shape[1])    
#    # If gain is okay, then score[i] = true


    score = np.abs(1 - gainShift) <= threshold

    # display results
    print("Target value at rate %3.3f is %3.3f" %(rateLook, binBaseline))
    print("Gain Shift Threshold =  %3.3f" % threshold)
    
    print("%s\t%s\t%s" %("Det#", "Shift", "Good?"))
    
    # detector # = index+1
    for (index, gs) in enumerate(gainShift):
        # Only show for not bad detectors
        if option == 'all':
                print("%d\t%3.3f\t%r" %(index+1, gs, score[index]))
        elif option == 'gooddet':
            if (index+1) not in badDetectors:
                print("%d\t%3.3f\t%r" %(index+1, gs, score[index]))
        elif option == 'badgain':
            if ~score[index]:
                print("%d\t%3.3f\t%r" %(index+1, gs, score[index]))
        elif option == 'badgaingooddet':
            if ~score[index] and (index+1) not in badDetectors:
                print("%d\t%3.3f\t%r" %(index+1, gs, score[index]))

if __name__ == "__main__":
    if len(sys.argv) < 5:
        CheckGain(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        CheckGain(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])
